#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
class KernelHistogram{
public:
    __aicore__ inline KernelHistogram(){}
    __aicore__ inline void Init(
        GM_ADDR x, GM_ADDR y, 
        int32_t totalLen, int32_t bins, float min, float max, float offset, float width, 
        int32_t blockLen, int32_t bufferLen, 
        int32_t groups, int32_t tailLen, int32_t tailLenPadded
    ){
        this->totalLen = totalLen;
        this->bins = bins;
        this->min = min;
        this->max = max;
        this->offset = offset;
        this->width = width;
        this->blockLen = blockLen;
        this->bufferLen = bufferLen;
        this->groups = groups;
        this->tailLen = tailLen;
        this->tailLenPadded = tailLenPadded;
        xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x, totalLen);
        yGm.SetGlobalBuffer((__gm__ DTYPE_X *)y, bins);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->bufferLen * sizeof(DTYPE_X));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->bufferLen * sizeof(DTYPE_X));
        pipe.InitBuffer(castBuf, this->bufferLen*sizeof(float));
        pipe.InitBuffer(zeroBuf, this->bufferLen*sizeof(float));
        pipe.InitBuffer(widthBuf, this->bufferLen*sizeof(float));
        pipe.InitBuffer(maskBuf, this->bufferLen*sizeof(float));
        pipe.InitBuffer(resBuf, this->bufferLen*sizeof(float));
        pipe.InitBuffer(sumBuf, this->bufferLen*sizeof(float));
        zeroLocal = zeroBuf.Get<float>();
        widthLocal = widthBuf.Get<float>();
        sumLocal = sumBuf.Get<float>();
        resLocal = resBuf.Get<float>();
        maskLocal = maskBuf.Get<uint8_t>();
        zeroLocalInt = zeroBuf.Get<int32_t>();
        widthLocalInt = widthBuf.Get<int32_t>();
    }
    __aicore__ inline void Process()
    {
        LocalTensor<int32_t> sumLocalInt = sumBuf.Get<int32_t>();
        for(int32_t i=0;i<bins;i++){
            left = min+(max-min)*static_cast<float>(i)/static_cast<float>(bins);
            right = min+(max-min)*static_cast<float>(i+1)/static_cast<float>(bins);
            Sub(sumLocalInt, sumLocalInt, sumLocalInt, bufferLen);
            Sub(zeroLocalInt, zeroLocalInt, zeroLocalInt, bufferLen);
            Sub(widthLocalInt, widthLocalInt, widthLocalInt, bufferLen);
            Adds(widthLocal, widthLocal, right, bufferLen);
            Adds(zeroLocal, zeroLocal, left, bufferLen);
            int32_t pos = 0;
            for(int32_t j=0;j<groups;j++){
                CopyIn(pos, bufferLen);
                Compute(bufferLen, i);
                pos+=bufferLen;
            }
            if(tailLen){
                CopyIn(pos, tailLenPadded);
                Compute(tailLen, i);
            }
            ComputeOut();
            CopyOut(i);
        }
    }
private:
    __aicore__ inline void CopyIn(int32_t pos, int32_t len)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopy(xLocal, xGm[pos], len);
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t len, int32_t idx)
    {
        // PRINTF("bin: %d left: %f, right: %f\n", idx, zeroLocal.GetValue(0), widthLocal.GetValue(0));
        LocalTensor<float>xLocal;
        LocalTensor<DTYPE_X> dx;
        if constexpr (std::is_same_v<DTYPE_X, float>){
            xLocal = inQueueX.DeQue<float>();

            // for(int32_t i=0;i<len;i++)PRINTF("%f ", xLocal.GetValue(i));PRINTF("\n");
            Sub(resLocal, resLocal, resLocal, len);
            Compare(maskLocal, xLocal, zeroLocal, CMPMODE::LT, len);
            Select(resLocal, maskLocal, resLocal, static_cast<float>(1), SELMODE::VSEL_TENSOR_SCALAR_MODE, len);
            // for(int32_t i=0;i<len;i++)PRINTF("%f ", resLocal.GetValue(i));PRINTF("\n");
            if(idx==bins-1) {
                Compare(maskLocal, xLocal, widthLocal, CMPMODE::LE, len);
            }
            else {
                Compare(maskLocal, xLocal, widthLocal, CMPMODE::LT, len);
            }
            Select(resLocal, maskLocal, resLocal, static_cast<float>(0), SELMODE::VSEL_TENSOR_SCALAR_MODE, len);
            // for(int32_t i=0;i<len;i++)PRINTF("%f ", resLocal.GetValue(i));PRINTF("\n");
            Add(sumLocal, resLocal, sumLocal, len);
            // for(int32_t i=0;i<len;i++)PRINTF("%f ", sumLocal.GetValue(i));PRINTF("\n");

            inQueueX.FreeTensor(xLocal);
        }
        else{
            xLocal = castBuf.Get<float>();
            dx = inQueueX.DeQue<DTYPE_X>();
            Cast(xLocal, dx, RoundMode::CAST_NONE, len);

            // for(int32_t i=0;i<len;i++)PRINTF("%f ", xLocal.GetValue(i));PRINTF("\n");
            Sub(resLocal, resLocal, resLocal, len);
            Compare(maskLocal, xLocal, zeroLocal, CMPMODE::LT, len);
            Select(resLocal, maskLocal, resLocal, static_cast<float>(1), SELMODE::VSEL_TENSOR_SCALAR_MODE, len);
            // for(int32_t i=0;i<len;i++)PRINTF("%f ", resLocal.GetValue(i));PRINTF("\n");
            if(idx==bins-1) {
                Compare(maskLocal, xLocal, widthLocal, CMPMODE::LE, len);
            }
            else {
                Compare(maskLocal, xLocal, widthLocal, CMPMODE::LT, len);
            }
            Select(resLocal, maskLocal, resLocal, static_cast<float>(0), SELMODE::VSEL_TENSOR_SCALAR_MODE, len);
            // for(int32_t i=0;i<len;i++)PRINTF("%f ", resLocal.GetValue(i));PRINTF("\n");
            Add(sumLocal, resLocal, sumLocal, len);
            // for(int32_t i=0;i<len;i++)PRINTF("%f ", sumLocal.GetValue(i));PRINTF("\n");
            inQueueX.FreeTensor(dx);
        }

    }
    __aicore__ inline void ComputeOut()
    {
        if constexpr (std::is_same_v<DTYPE_X, float>){
            LocalTensor<float> yLocalFloat = outQueueY.AllocTensor<float>();
            ReduceSum(yLocalFloat, sumLocal, sumLocal, bufferLen);
            outQueueY.EnQue(yLocalFloat);
        }
        else{
            LocalTensor<DTYPE_X> yLocal = outQueueY.AllocTensor<DTYPE_X>();
            ReduceSum(sumLocal, sumLocal, sumLocal, bufferLen);
            Cast(yLocal, sumLocal, RoundMode::CAST_RINT, 1);
            outQueueY.EnQue(yLocal);
        }
        
        
    }
    __aicore__ inline void CopyOut(int32_t idx)
    {
        LocalTensor<DTYPE_X> yLocal = outQueueY.DeQue<DTYPE_X>();
        if(idx+blockLen<=bins){
            DataCopy(yGm[idx], yLocal, blockLen);
        }
        else{
            yGm.SetValue(idx, yLocal.GetValue(0));
        }
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<TPosition::VECCALC> castBuf, zeroBuf, widthBuf, maskBuf, resBuf, sumBuf;
    GlobalTensor<DTYPE_X> xGm, yGm;
    int32_t totalLen, bins, blockLen, bufferLen, groups, tailLen, tailLenPadded;
    float offset, binsf, width, min, max, left, right;
    LocalTensor<float> zeroLocal, widthLocal, sumLocal, resLocal;
    LocalTensor<uint8_t> maskLocal;
    LocalTensor<int32_t> zeroLocalInt, widthLocalInt;
};

extern "C" __global__ __aicore__ void histogram(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelHistogram op;
    op.Init(
        x, y, 
        tiling_data.totalLen, tiling_data.bins, tiling_data.min, tiling_data.max, tiling_data.offset, tiling_data.width, 
        tiling_data.blockLen, tiling_data.bufferLen, 
        tiling_data.groups, tiling_data.tailLen, tiling_data.tailLenPadded
    );
    op.Process();
}